"""
Author: Morphlng
Date: 2024-03-19 14:00:00
LastEditTime: 2024-03-25 10:34:05
LastEditors: Morphlng
Description: Test for transform
FilePath: /DrivingGym/tests/common/geom/transform_test.py
"""

import math
import random

import numpy as np

from driving_gym.simulation.common.geom import Location, Rotation, Transform
from driving_gym.simulation.common.stub import try_import_carla

carla = try_import_carla()


class TestTransform:

    def test_initialization_with_location_and_rotation(self):
        loc = Location(1, 2, 3)
        rot = Rotation(45, 45, 45)
        trans = Transform(location=loc, rotation=rot)
        assert np.array_equal(
            trans.location.as_numpy_array(), np.array([1, 2, 3])
        ), "Location initialization failed."
        assert np.array_equal(
            trans.rotation.as_numpy_array(), np.array([45, 45, 45])
        ), "Rotation initialization failed."

        carla_loc = carla.Location(1, 2, 3)
        carla_rot = carla.Rotation(45, 45, 45)
        carla_trans = carla.Transform(carla_loc, carla_rot)
        carla_mat = carla_trans.get_matrix()
        assert np.allclose(
            trans.matrix, np.array(carla_mat), atol=1e-5
        ), "Matrix incorrect."

    def test_initialization_with_matrix(self):
        matrix = np.array([[1, 0, 0, 1], [0, 1, 0, 2], [0, 0, 1, 3], [0, 0, 0, 1]])
        trans = Transform(matrix=matrix)
        assert np.allclose(
            trans.location.as_numpy_array(), np.array([1, 2, 3])
        ), "Location from matrix failed."
        assert trans.matrix is not None, "Matrix initialization failed."

        carla_loc = carla.Location(1, 2, 3)
        carla_rot = carla.Rotation(0, 0, 0)
        carla_trans = carla.Transform(carla_loc, carla_rot)
        carla_mat = carla_trans.get_matrix()
        assert np.allclose(
            matrix, np.array(carla_mat), atol=1e-5
        ), "Transform incorrect."

    def test_get_matrix(self):
        trans = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        carla_trans = carla.Transform(
            carla.Location(1, 2, 3), carla.Rotation(45, 45, 45)
        )
        carla_mat = carla_trans.get_matrix()
        assert np.allclose(
            trans.matrix, np.array(carla_mat), atol=1e-5
        ), "Matrix computation failed."

    def test_get_inverse_matrix(self):
        trans = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        inv_matrix = trans.get_inverse_matrix()

        carla_trans = carla.Transform(
            carla.Location(1, 2, 3), carla.Rotation(45, 45, 45)
        )
        carla_inv_matrix = carla_trans.get_inverse_matrix()

        assert np.allclose(
            inv_matrix, np.array(carla_inv_matrix), atol=1e-5
        ), "Inverse matrix computation failed."

    def test_transform_point(self):
        for _ in range(100):
            loc = Location(
                random.uniform(-100, 100),
                random.uniform(-100, 100),
                random.uniform(-100, 100),
            )
            rot = Rotation(
                random.uniform(-180, 180),
                random.uniform(-180, 180),
                random.uniform(-180, 180),
            )
            carla_loc = carla.Location(loc.x, loc.y, loc.z)
            carla_rot = carla.Rotation(rot.pitch, rot.yaw, rot.roll)
            point = (
                random.uniform(-100, 100),
                random.uniform(-100, 100),
                random.uniform(-100, 100),
            )
            carla_point = carla.Location(*point)

            res = Transform(loc, rot).transform_point(np.array(point))
            carla_res = carla.Transform(carla_loc, carla_rot).transform(carla_point)

            assert np.allclose(
                res, np.array([carla_res.x, carla_res.y, carla_res.z]), atol=1e-3
            ), "Point transformation incorrect."

    def test_inverse_transform_point(self):
        trans = Transform(location=Location(1, 2, 3), rotation=Rotation(0, 0, 0))
        point = np.array([1, 1, 1])
        inv_transformed_point = trans.inverse_transform_point(point)
        assert np.allclose(
            inv_transformed_point, np.array([0, -1, -2])
        ), "Inverse point transformation failed."

    def test_transform_location(self):
        for _ in range(100):
            loc = Location(
                random.uniform(-100, 100),
                random.uniform(-100, 100),
                random.uniform(-100, 100),
            )
            rot = Rotation(
                random.uniform(-180, 180),
                random.uniform(-180, 180),
                random.uniform(-180, 180),
            )
            carla_loc = carla.Location(loc.x, loc.y, loc.z)
            carla_rot = carla.Rotation(rot.pitch, rot.yaw, rot.roll)
            point = (
                random.uniform(-100, 100),
                random.uniform(-100, 100),
                random.uniform(-100, 100),
            )
            location = Location(*point)
            carla_location = carla.Location(*point)

            res = Transform(loc, rot).transform_location(location)
            carla_res = carla.Transform(carla_loc, carla_rot).transform(carla_location)

            assert np.allclose(
                res.as_numpy_array(),
                np.array([carla_res.x, carla_res.y, carla_res.z]),
                atol=1e-3,
            ), "Location transformation incorrect."

    def test_inverse_transform_location(self):
        trans = Transform(location=Location(1, 2, 3), rotation=Rotation(0, 0, 0))
        loc = Location(1, 1, 1)
        inv_transformed_loc = trans.inverse_transform_location(loc)
        assert (
            isinstance(inv_transformed_loc, Location)
            and inv_transformed_loc.x == 0
            and inv_transformed_loc.y == -1
            and inv_transformed_loc.z == -2
        ), "Inverse location transformation failed."

    def test_transformations_composition(self):
        trans1 = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        trans2 = Transform(location=Location(4, 5, 6), rotation=Rotation(30, 60, 90))
        combined_trans = trans1 * trans2
        assert (
            combined_trans.location != trans1.location
            and combined_trans.rotation != trans1.rotation
        ), "Transformations composition failed."

    def test_equality(self):
        trans1 = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        trans2 = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        assert trans1 == trans2, "Transform equality check failed."

    def test_inequality(self):
        trans1 = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        trans2 = Transform(location=Location(4, 5, 6), rotation=Rotation(30, 60, 90))
        assert trans1 != trans2, "Transform inequality check failed."

    def test_get_angle_and_magnitude(self):
        trans = Transform(location=Location(0, 0, 0), rotation=Rotation(0, 0, 0))
        target_loc = Location(10, 0, 0)
        angle, magnitude = trans.get_angle_and_magnitude(target_loc)
        assert math.isclose(angle, 0, abs_tol=1e-5), "Angle computation failed."
        assert math.isclose(
            magnitude, 10, abs_tol=1e-5
        ), "Magnitude computation failed."

    def test_is_within_distance_ahead(self):
        trans = Transform(location=Location(0, 0, 0), rotation=Rotation(0, 0, 0))
        target_loc = Location(10, 0, 0)  # Directly ahead
        assert trans.is_within_distance_ahead(
            target_loc, 15
        ), "Target should be within distance ahead."
        assert not trans.is_within_distance_ahead(
            target_loc, 5
        ), "Target should not be within distance ahead."

        target_loc_side = Location(0, 10, 0)  # To the side
        assert not trans.is_within_distance_ahead(
            target_loc_side, 15
        ), "Target should not be within distance ahead (side)."

    def test_inverse_transform(self):
        trans = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        inv_trans = trans.inverse_transform()
        original_point = np.array([1, 1, 1])
        transformed_point = trans.transform_point(original_point)
        inv_transformed_point = inv_trans.transform_point(transformed_point)
        assert np.allclose(
            inv_transformed_point, original_point
        ), "Inverse transform failed."

    def test_mul(self):
        trans1 = Transform(location=Location(1, 2, 3), rotation=Rotation(45, 45, 45))
        trans2 = Transform(location=Location(4, 5, 6), rotation=Rotation(30, 60, 90))
        result_trans = trans1 * trans2
        assert isinstance(
            result_trans, Transform
        ), "Multiplication failed to return a Transform object."
